import requests
import psycopg2
from datetime import datetime
from psycopg2 import sql
from typing import Dict, List, Optional, Any

# 获取accessToken
def get_access_token(app_key: str, app_secret: str) -> str:
    """获取钉钉访问令牌"""
    url = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
    payload = {
        'appKey': app_key,
        'appSecret': app_secret
    }
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json().get('accessToken')
    except requests.exceptions.RequestException as e:
        raise Exception(f"获取访问令牌失败: {e}")

# 获取企业考勤排班详情
def fetch_listschedule(access_token: str, work_date: str) -> List[Dict[str, Any]]:
    """获取指定日期的考勤排班信息"""
    url = "https://oapi.dingtalk.com/topapi/attendance/listschedule"
    headers = {
        'Content-Type': 'application/json'
    }
    params = {
        'access_token': access_token
    }
    body = {
        'workDate': work_date
    }
    try:
        response = requests.post(url, headers=headers, params=params, json=body)
        response.raise_for_status()
        data = response.json()
        if data.get('errcode') == 0:
            result = data.get('result', {})
            return result.get('schedules', [])
        return []
    except requests.exceptions.RequestException as e:
        print(f"获取排班信息失败: {e}")
        return []

# 获取计划考勤时间
def get_plan_check_time(schedules: List[Dict[str, Any]], userid: str, check_type: str) -> Optional[str]:
    """获取指定用户的计划考勤时间"""
    for schedule in schedules:
        if schedule.get('userid') == userid and schedule.get('check_type') == check_type:
            return schedule.get('plan_check_time')
    return None

# 获取所有排班信息
def get_all_schedules(access_token: str, work_date: str) -> List[Dict[str, Any]]:
    """获取并处理所有排班信息"""
    schedules = fetch_listschedule(access_token, work_date)
    user_schedules = {}

    for schedule in schedules:
        userid = schedule.get('userid')
        if not userid:
            continue

        if userid not in user_schedules:
            user_schedules[userid] = {
                'userid': userid,
                'class_id': schedule.get('class_id',None),
                'check_on_duty_time': get_plan_check_time(schedules, userid, 'OnDuty'),
                'check_off_duty_time': get_plan_check_time(schedules, userid, 'OffDuty'),
                'work_date': work_date
            }

    return list(user_schedules.values())

# 查询当前用户数据是否存在
def exists_schedules(schedules: List[Dict[str, Any]], userid: str) -> bool:
    """检查用户排班数据是否存在"""
    return any(schedule.get('userid') == userid for schedule in schedules)

# 连接到数据库并插入数据
def insert_schedule_data(schedules: List[Dict[str, Any]], work_date: str) -> None:
    """将排班数据插入数据库"""
    if not schedules:
        print("没有排班数据需要插入")
        return

    conn = None
    try:
        # 从环境变量或配置文件获取数据库连接信息
        conn = psycopg2.connect(
            dbname="metabase",
            user="postgres",
            password="postgres",
            host="172.16.5.226",
            port="32685"
        )

        with conn.cursor() as cur:
            # 清空当前天下的数据
            cur.execute(
                "DELETE FROM tp.project_member_shift_record WHERE work_date::date = %s",
                (work_date,)
            )

            # 批量插入数据
            insert_query = sql.SQL("""
                                   INSERT INTO tp.project_member_shift_record
                                       (userid, check_on_duty_time, check_off_duty_time, work_date,data_source,class_id)
                                   VALUES %s
                                   """)

            # 准备数据
            data = [
                (
                    schedule['userid'],
                    schedule['check_on_duty_time'],
                    schedule['check_off_duty_time'],
                    f"{work_date} 00:00:00",
                    'dingding',
                    schedule['class_id']
                )
                for schedule in schedules
            ]

            # 使用execute_values进行批量插入
            from psycopg2.extras import execute_values
            execute_values(cur, insert_query, data)

        conn.commit()
        print(f"成功插入 {len(schedules)} 条排班数据")

    except psycopg2.Error as e:
        print(f"数据库操作失败: {e}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()

# 主函数
def main():
    """主程序入口"""
    try:
        # 获取access_token
        app_key = 'dingtyw4aiyucyykv1rr'
        app_secret = 'fkGxuNZ653SH_CyAdT7mhSCfHy40idkCAUgc_6bSk_QceBjBfg79taOBmVnLGyPO'
        access_token = get_access_token(app_key, app_secret)

        if not access_token:
            raise Exception("获取访问令牌失败")

        # 获取当前日期，格式为 YYYY-MM-DD
        work_date = datetime.now().strftime("%Y-%m-%d")

        # 获取企业考勤排班详情
        schedules = get_all_schedules(access_token, work_date)

        if not schedules:
            print("当前无排班数据")
        else:
            insert_schedule_data(schedules, work_date)
            print("排班数据已成功插入数据库")

    except Exception as e:
        print(f"程序执行出错: {e}")

if __name__ == "__main__":
    main()